import numpy as np


class Shapefunc:
    def __init__(self, mnode, ndim):
        self.mnode = mnode
        self.ndim = ndim

    def fun(self, ks, yt, zita):
        N = np.zeros(self.mnode)
        if self.ndim == 2:
            if self.mnode == 4:
                N[0] = (1 - ks) * (1 - yt) / 4
                N[1] = (1 + ks) * (1 - yt) / 4
                N[2] = (1 + ks) * (1 + yt) / 4
                N[3] = (1 - ks) * (1 + yt) / 4
            elif self.mnode == 8:
                N[0] = (1 / 4) * (1 + (-1) * ks) * (1 + (-1) * yt) * ((-1) + (-1) * ks + (-1) * yt)
                N[1] = (1 / 4) * (1 + ks) * (1 + (-1) * yt) * ((-1) + ks + (-1) * yt)
                N[2] = (1 / 4) * (1 + ks) * (1 + yt) * ((-1) + ks + yt)
                N[3] = (1 / 4) * (1 + (-1) * ks) * (1 + yt) * ((-1) + (-1) * ks + yt)
                N[4] = (1 / 2) * (1 + (-1) * ks ** 2) * (1 + (-1) * yt)
                N[5] = (1 / 2) * (1 + ks) * (1 + (-1) * yt ** 2)
                N[6] = (1 / 2) * (1 + (-1) * ks ** 2) * (1 + yt)
                N[7] = (1 / 2) * (1 + (-1) * ks) * (1 + (-1) * yt ** 2)
        else:
            if self.mnode == 8:
                N[0] = (1 / 8) * (1 + ks) * (1 + (-1) * yt) * (1 + (-1) * zita)
                N[1] = (1 / 8) * (1 + ks) * (1 + yt) * (1 + (-1) * zita)
                N[2] = (1 / 8) * (1 + (-1) * ks) * (1 + yt) * (1 + (-1) * zita)
                N[3] = (1 / 8) * (1 + (-1) * ks) * (1 + (-1) * yt) * (1 + (-1) * zita)
                N[4] = (1 / 8) * (1 + ks) * (1 + (-1) * yt) * (1 + zita)
                N[5] = (1 / 8) * (1 + ks) * (1 + yt) * (1 + zita)
                N[6] = (1 / 8) * (1 + (-1) * ks) * (1 + yt) * (1 + zita)
                N[7] = (1 / 8) * (1 + (-1) * ks) * (1 + (-1) * yt) * (1 + zita)
            elif self.mnode == 20:
                N[0] = (1 / 8) * (1 + ks) * (1 + (-1) * yt) * (1 + (-1) * zita) * ((-2) + ks + (-1) * yt + (
                    -1) * zita)
                N[1] = (1 / 8) * (1 + ks) * (1 + yt) * (1 + (-1) * zita) * ((-2) + ks + yt + (-1) * zita)
                N[2] = (1 / 8) * (1 + (-1) * ks) * (1 + yt) * (1 + (-1) * zita) * ((-2) + (-1) * ks + yt + (
                    -1) * zita)
                N[3] = (1 / 8) * (1 + (-1) * ks) * (1 + (-1) * yt) * (1 + (-1) * zita) * ((-2) + (-1) * ks +
                                                                                          (-1) * yt + (-1) * zita)
                N[4] = (1 / 8) * (1 + ks) * (1 + (-1) * yt) * (1 + zita) * ((-2) + ks + (-1) * yt + zita)
                N[5] = (1 / 8) * (1 + ks) * (1 + yt) * (1 + zita) * ((-2) + ks + yt + zita)
                N[6] = (1 / 8) * (1 + (-1) * ks) * (1 + yt) * (1 + zita) * ((-2) + (-1) * ks + yt + zita)
                N[7] = (1 / 8) * (1 + (-1) * ks) * (1 + (-1) * yt) * (1 + zita) * ((-2) + (-1) * ks + (-1) *
                                                                                   yt + zita)
                N[8] = (1 / 4) * (1 + ks) * (1 + (-1) * yt ** 2) * (1 + (-1) * zita)
                N[9] = (1 / 4) * (1 + (-1) * ks ** 2) * (1 + yt) * (1 + (-1) * zita)
                N[10] = (1 / 4) * (1 + (-1) * ks) * (1 + (-1) * yt ** 2) * (1 + (-1) * zita)
                N[11] = (1 / 4) * (1 + (-1) * ks ** 2) * (1 + (-1) * yt) * (1 + (-1) * zita)
                N[12] = (1 / 4) * (1 + ks) * (1 + (-1) * yt ** 2) * (1 + zita)
                N[13] = (1 / 4) * (1 + (-1) * ks ** 2) * (1 + yt) * (1 + zita)
                N[14] = (1 / 4) * (1 + (-1) * ks) * (1 + (-1) * yt ** 2) * (1 + zita)
                N[15] = (1 / 4) * (1 + (-1) * ks ** 2) * (1 + (-1) * yt) * (1 + zita)
                N[16] = (1 / 4) * (1 + ks) * (1 + (-1) * yt) * (1 + (-1) * zita ** 2)
                N[17] = (1 / 4) * (1 + ks) * (1 + yt) * (1 + (-1) * zita ** 2)
                N[18] = (1 / 4) * (1 + (-1) * ks) * (1 + yt) * (1 + (-1) * zita ** 2)
                N[19] = (1 / 4) * (1 + (-1) * ks) * (1 + (-1) * yt) * (1 + (-1) * zita ** 2)

        return N
